package com.ruoyi.system.websocket;

import cn.hutool.json.JSONObject;
import cn.hutool.json.JSONUtil;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;

import javax.annotation.PostConstruct;
import javax.annotation.PreDestroy;
import javax.websocket.*;
import javax.websocket.server.PathParam;
import javax.websocket.server.ServerEndpoint;
import java.io.IOException;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CopyOnWriteArraySet;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;

@ServerEndpoint(value = "/ws/asset")
@Component
@Slf4j
public class WebSocketServer {

    private static final AtomicInteger OnlineCount = new AtomicInteger(0);
    private static CopyOnWriteArraySet<Session> SessionSet = new CopyOnWriteArraySet<>();
    // 存储token和session的映射关系
    private static final Map<String, Session> tokenSessionMap = new ConcurrentHashMap<>();
    // 存储session和token的映射关系
    private static final Map<Session, String> sessionTokenMap = new ConcurrentHashMap<>();

    // 心跳间隔(秒)
    private static final long HEARTBEAT_INTERVAL = 30;
    // 超时时间(秒)
    private static final long TIMEOUT = 65;

    // 存储每个session的最后活跃时间
    private static final Map<String, Long> lastActiveTime = new ConcurrentHashMap<>();

    // 心跳检测线程池
    private static final ScheduledExecutorService scheduler = Executors.newSingleThreadScheduledExecutor();

    @PostConstruct
    public void init() {
        // 启动心跳检测任务
        scheduler.scheduleAtFixedRate(this::checkHeartbeats,
                HEARTBEAT_INTERVAL,
                HEARTBEAT_INTERVAL,
                TimeUnit.SECONDS);
    }

    @PreDestroy
    public void destroy() {
        scheduler.shutdown();
    }

    /**
     * 发送消息
     */
    public static void SendMessage(Session session, String message) {
        try {
            session.getBasicRemote().sendText(message);
            log.info("socket发送消息: {}", message);
        } catch (Exception e) {
            log.error("发送消息出错：{}", e.getMessage());
            e.printStackTrace();
        }
    }

    /**
     * 根据token发送消息
     */
    public static void sendMessageByToken(String token, String message) {
        Session session = tokenSessionMap.get(token);
        if (session != null && session.isOpen()) {
            SendMessage(session, message);
        } else {
            log.warn("用户token={}的WebSocket连接不存在或已关闭", token);
        }
    }

    @OnOpen
    public void onOpen(Session session) {
        SessionSet.add(session);
        int cnt = OnlineCount.incrementAndGet();
        lastActiveTime.put(session.getId(), System.currentTimeMillis());
        log.info("有连接加入，当前连接数为：{}", cnt);

        // 发送欢迎消息
        SendMessage(session, "{\"type\":\"welcome\",\"message\":\"Connected. Server will send heartbeat every " + HEARTBEAT_INTERVAL + " seconds.\"}");
    }

    @OnClose
    public void onClose(Session session, @PathParam("cid") String cid) {
        SessionSet.remove(session);
        int cnt = OnlineCount.decrementAndGet();
        lastActiveTime.remove(session.getId());
        
        // 清理token映射关系
        String token = sessionTokenMap.remove(session);
        if (token != null) {
            tokenSessionMap.remove(token);
            log.info("清理token映射关系, token={}, sessionId={}", token, session.getId());
        }
        
        log.info("有连接关闭，当前连接数为：{}, CID: {}", cnt, cid);
    }

    @OnMessage
    public void onMessage(String message, Session session) {
        try {
            // 更新最后活跃时间
            lastActiveTime.put(session.getId(), System.currentTimeMillis());

            JSONObject messageObject = JSONUtil.parseObj(message);
            log.info("收到消息: {}", messageObject.toString());
            
            // 处理心跳回复
            if ("heartbeat".equals(messageObject.getStr("type"))) {
                String token = messageObject.getStr("token");
                if (token != null) {
                    // 清理旧的映射关系
                    String oldToken = sessionTokenMap.get(session);
                    if (oldToken != null && !oldToken.equals(token)) {
                        tokenSessionMap.remove(oldToken);
                        log.info("清理旧的token映射关系, oldToken={}, sessionId={}", oldToken, session.getId());
                    }
                    
                    // 更新token和session的映射关系
                    tokenSessionMap.put(token, session);
                    sessionTokenMap.put(session, token);
                    log.info("更新token映射关系, token={}, sessionId={}", token, session.getId());
                }
                log.debug("收到心跳回复: {}", session.getId());
                SendMessage(session, "{\"type\":\"heartbeat\",\"message\":\"heartbeat\"}");
                return;
            }

            // 处理其他业务消息
            SendMessage(session, "{\"orderSn\":\"123456\"}");

        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    @OnError
    public void onError(Session session, Throwable error) {
        log.error("发生错误：{}，Session ID： {}", error.getMessage(), session.getId());
        error.printStackTrace();
    }

    /**
     * 心跳检测方法
     */
    private void checkHeartbeats() {
        long currentTime = System.currentTimeMillis();

        for (Session session : SessionSet) {
            String sessionId = session.getId();
            Long lastActive = lastActiveTime.get(sessionId);

            if (lastActive == null || currentTime - lastActive > TIMEOUT * 1000) {
                // 超时，关闭连接
                try {
                    session.close(new CloseReason(CloseReason.CloseCodes.NORMAL_CLOSURE, "Connection timeout"));
                    log.info("连接超时关闭: {}", sessionId);
                } catch (IOException e) {
                    log.error("关闭连接出错: {}", e.getMessage());
                }
                SessionSet.remove(session);
                OnlineCount.decrementAndGet();
                lastActiveTime.remove(sessionId);
                
                // 清理token映射关系
                String token = sessionTokenMap.remove(session);
                if (token != null) {
                    tokenSessionMap.remove(token);
                    log.info("清理超时连接的token映射关系, token={}, sessionId={}", token, sessionId);
                }
            } else if (currentTime - lastActive > HEARTBEAT_INTERVAL * 1000) {
                // 发送心跳
                try {
                    session.getBasicRemote().sendText("{\"type\":\"heartbeat\"}");
                    log.debug("发送心跳到: {}", sessionId);
                } catch (IOException e) {
                    log.error("发送心跳失败: {}", e.getMessage());
                    SessionSet.remove(session);
                    OnlineCount.decrementAndGet();
                    lastActiveTime.remove(sessionId);
                    
                    // 清理token映射关系
                    String token = sessionTokenMap.remove(session);
                    if (token != null) {
                        tokenSessionMap.remove(token);
                        log.info("清理心跳失败的token映射关系, token={}, sessionId={}", token, sessionId);
                    }
                }
            }
        }
    }

    /**
     * 获取所有WebSocket会话
     */
    public static CopyOnWriteArraySet<Session> getSessionSet() {
        return SessionSet;
    }
}